import asyncio
import websockets
import json
import cv2
import base64
import numpy as np
import os

async def test_websocket():
    uri = "ws://202.114.212.194:50511/detect"  # your server addr
    cap = cv2.VideoCapture(0)

    # 创建保存图片的目录
    save_dir = "received_images"
    os.makedirs(save_dir, exist_ok=True)
    
    try:
        async with websockets.connect(uri) as websocket:
            frame_count = 0
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    print("Failed to read frame")
                    break

                # 编码图片为 Base64
                _, buffer = cv2.imencode('.jpg', frame)
                img_str = base64.b64encode(buffer)

                # 发送图片数据给服务器
                await websocket.send(img_str)
                # print("Sent frame to server")

                # 接收服务器返回的推理结果
                response = await websocket.recv()
                if isinstance(response, str):
                    result = json.loads(response)
                    print(f'Received: {result["detected_class"]}')
                elif isinstance(response, bytes):
                    # 解码接收到的图片数据
                    np_img = np.frombuffer(response, dtype=np.uint8)
                    img = cv2.imdecode(np_img, cv2.IMREAD_COLOR)

                    # 保存图片到本地文件
                    img_filename = os.path.join(save_dir, f"frame_{frame_count:04d}.jpg")
                    # cv2.imwrite(img_filename, img)
                    print(f"Saved {img_filename}")
            

                frame_count += 1

    except Exception as e:
        print(f"Failed to connect: {e}")

    cap.release()

# run
asyncio.get_event_loop().run_until_complete(test_websocket())